import json
import os
import uuid

from src.custom.tools import logger
from dotenv import load_dotenv
from openai import OpenAI
from zhipuai import ZhipuAI
from src.custom.agent import ChatAI
from src.custom.messages import HistoryLoader, build_human_message, build_ai_message, HumanMessage, build_system_message
from src.custom.tools import ToolDeclarationConfig, write_json, read_json
from src.user import UserInfo

tool = ToolDeclarationConfig("custom.tools")

@tool(location="要查询的城市名", format="使用的温度单位")
def get_current_weather(location: str, format: str):
    """获取一个地方的天气信息"""
    return "当前天气：报告天气：晴朗：0°C，晴朗，空气质量：15，风速：东北风 2级，湿度：44%，体感温度：0°，能见度：30 公里，紫外线指数：0，气压：1028 hPa，露点：-11°"

@tool(input_str="用户想要搜索的内容")
def baidu(input_str: str):
    """通过百度搜索用户的内容"""
    print("百度中 ... "+input_str)
    return input_str+"未能找到"

tools = [get_current_weather, baidu]

class ZhiPuAI(ChatAI):
    def _get_llm_client(self):
        client = ZhipuAI(api_key=os.getenv("ZHIPU_API_KEY"))
        model_name = 'GLM-4-Flash'
        return model_name, client
class AliAI(ChatAI):
    def _get_llm_client(self):
        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        model_name = 'qwen-long'
        return model_name, client

# 历史加载器需要管理用户账号的操作，AI需要管理加载器的操作，不可管理用户操作
class MyHistoryLoader(HistoryLoader):
    def __init__(self, user_acc: UserInfo, system_prompt: str = None):
        self.history_pointer = 0
        self.system_prompt = build_system_message(system_prompt if system_prompt else "")    # 永久提示词，新开启会话的时候总会存在
        self.user_acc = user_acc                                    # 用户信息，管理历史加载器加载的历史
        if self.user_acc.session_name is None:
            new_session = str(uuid.uuid4())                         # TODO 可能需要改一下初始名，或者看需不需要保存
            self.user_acc.new_session(new_session)
            self.user_acc.change_session(new_session)
        self.session_storage = ""                                   # 因为跟着session变，所以在必要的时候需要更新一下值
        self.files_storage = ""                                     # 文件存储的目录
        self.files_content_ = ""
        self.dir_work = os.path.join(".", "storage", user_acc.name)
        self.history = None                                         # 等需要历史的时候再去加载，可以延长之后配置历史的时机

    def renewal_session_storage(self):
        """更新self.session_storage的存储路径"""
        self.session_storage = os.path.join(
            self.dir_work,
            "sessions",
            self.user_acc.session_name + ".json"
        )

    def renewal_files_storage(self):
        self.files_storage = os.path.join(
            self.dir_work,
            "files",
        )

    @tool(prefix="json.storage")
    def get_history(self) -> list:
        """获取聊天记录"""
        if self.history is None:
            self.history = self.load()
        return self.history

    @tool(prefix="json.storage")
    def save(self):
        """可以保存当前对话的聊天历史"""
        self.renewal_session_storage()
        deleted_history = []

        for i in range(self.history_pointer, len(self.history)):
            message = self.history[i]
            if 'is_abandoned' in message and message['is_abandoned'] == 'true':
                deleted_history.append(message)
        for message in deleted_history:
            self.history.remove(message)
        self.history_pointer = len(self.history)
        if write_json(self.session_storage, self.history):
            # TODO 是否在这里保存
            self.user_acc.save()
            return "保存成功"
        else:
            return "保存失败"

    @tool(prefix="json.storage")
    def load(self):
        """可以根据用户的信息加载对应会话的历史记录"""
        self.renewal_session_storage()
        flag = read_json(self.session_storage)
        self.history = flag if flag else [self.system_prompt]
        self.history_pointer = len(self.history)
        return self.history

    @tool(prefix="json.storage", session_id="新的会话主题名")
    def set_user_session_name(self, session_id: str):
        """设置用户当前的会话标题，当用户没有指定的时候根据对话内容自定义，如果指定则设置为指定的名字"""
        old_session_name = self.user_acc.session_name
        new_session_path_name = os.path.join(
            self.dir_work,
            session_id + ".json"
        )
        # 重命名
        if os.path.exists(self.session_storage):
            os.rename(self.session_storage, new_session_path_name)
        self.user_acc.set_session_name(session_id)
        # 更新storage路径
        self.renewal_session_storage()
        return "会话名已经从《"+old_session_name+"》修改为：《"+self.user_acc.session_name+"》"

    @tool(prefix="json.storage")
    def get_current_session(self):
        """获取当前的会话主题的名字"""
        return self.user_acc.session_name

    @tool(prefix="json.storage", session_id="切换到叫这个名字的会话")
    def change_session(self, session_id: str):
        """切换、更改当前的会话主题"""
        # 先保存当前会话历史，再切换，加载新会话的历史
        self.save()
        if self.user_acc.change_session(session_id):
            self.history = self.load()
            return "《" + session_id + "》会话切换成功"
        else:
            return f"《{session_id}》会话切换失败，可能是不存在"

    @tool(prefix="json.storage", session_id="开启新会话的名字")
    def start_new_session(self, session_id: str) -> str:
        """开启一个新会话，如果没有指定会话主题先随机一个名字，之后根据会话内容再重新设置主题名"""
        if self.user_acc.new_session(session_id):
            self.change_session(session_id)
            return "开启了一个新的会话："+session_id
        else:
            return "开启新会话失败，可能此会话已经存在"

    @tool(prefix="json.storage")
    def get_session_list(self) -> list:
        """获取用户所有会话的列表，得到一个数组，每个元素都是会话的标题"""
        session_list = self.user_acc.session_list
        self.user_acc.session_list = [s for s in session_list if os.path.exists(os.path.join(self.dir_work, s+".json"))]
        # 已经保存的会话和当前会话，去除掉没有保存的空会话
        # TODO 逻辑不够好，处理没有保存的会话名，从list中剔除
        self.user_acc.add_session_to_list(self.user_acc.session_name)
        return self.user_acc.session_list

    # 这里可以过滤选择哪些消息该加入
    def add_message(self, message: dict[str, str]):
        # TODO 通过一个窗口指针来统一处理不用的信息
        self.history.append(message)

    def add_messages(self, messages: list[dict[str, str]]):
        for message in messages:
            self.add_message(message)

if __name__ == "__main__":
    load_dotenv("../.env")
    zhipu = ZhiPuAI()
    print(zhipu.send_message("下午好"))

    # user = UserInfo("zhangsan")
    # user.load()
    # sys_prompt = f"我叫：{user.name}，你是我的资深管家，管理着与我沟通的许多会话记录，每个会话都有一个主题（会话）名，尽你可能的功能帮助我解决问题。"
    # my_history_loader = MyHistoryLoader(user, sys_prompt)
    #
    # zhipu = ZhiPuAI(my_history_loader)
    # zhipu.bind_tools(tools)
    #
    # while True:
    #     try:
    #         user_input = input(f"\n\033[31m[{my_history_loader.user_acc.session_name}]\033[0m User:\n")
    #         if user_input in ['q', 'quit', 'exit', '退出']:
    #             break
    #         re = zhipu.send_message_with_tools(build_human_message(user_input))
    #         print(f"\n\033[31m[{my_history_loader.user_acc.session_name}]\033[0m \033[32mAI:\033[0m\n{re}")
    #     except Exception as e:
    #         raise e